package bma.common.langutil.concurrent;

import java.util.ArrayList;
import java.util.ConcurrentModificationException;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.AbstractExecutorService;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * see java.util.concurrent.ThreadPoolExecutor
 * 
 * @author guanzhong
 * 
 */
public class BMAThreadPoolExecutor extends AbstractExecutorService {

	final org.slf4j.Logger log = org.slf4j.LoggerFactory
			.getLogger(BMAThreadPoolExecutor.class);

	private static final RuntimePermission shutdownPerm = new RuntimePermission(
			"modifyThread");

	volatile int runState;
	static final int RUNNING = 0;
	static final int SHUTDOWN = 1;
	static final int STOP = 2;
	static final int TERMINATED = 3;

	private final BlockingQueue<Runnable> workQueue;

	private final ReentrantLock mainLock = new ReentrantLock();
	private final Condition termination = mainLock.newCondition();
	private final HashSet<Worker> workers = new HashSet<Worker>();
	private volatile long keepAliveTime;

	private volatile boolean allowCoreThreadTimeOut;

	private volatile int corePoolSize;

	private volatile int maximumPoolSize;

	private volatile int poolSize;

	private volatile BMARejectedExecutionHandler handler;

	private volatile ThreadFactory threadFactory;

	private int largestPoolSize;

	private AtomicLong completedTaskCount = new AtomicLong();
	private AtomicLong taskCount = new AtomicLong();

	private AtomicInteger waitingWorker = new AtomicInteger();

	private static final BMARejectedExecutionHandler defaultHandler = new AbortPolicy();

	// Constructors
	public BMAThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
			long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
		this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
				Executors.defaultThreadFactory(), defaultHandler);
	}

	public BMAThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
			long keepAliveTime, TimeUnit unit,
			BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
		this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
				threadFactory, defaultHandler);
	}

	public BMAThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
			long keepAliveTime, TimeUnit unit,
			BlockingQueue<Runnable> workQueue,
			BMARejectedExecutionHandler handler) {
		this(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
				Executors.defaultThreadFactory(), handler);
	}

	public BMAThreadPoolExecutor(int corePoolSize, int maximumPoolSize,
			long keepAliveTime, TimeUnit unit,
			BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory,
			BMARejectedExecutionHandler handler) {
		if (corePoolSize < 0 || maximumPoolSize <= 0
				|| maximumPoolSize < corePoolSize || keepAliveTime < 0)
			throw new IllegalArgumentException();
		if (workQueue == null || threadFactory == null || handler == null)
			throw new NullPointerException();
		this.corePoolSize = corePoolSize;
		this.maximumPoolSize = maximumPoolSize;
		this.workQueue = workQueue;
		this.keepAliveTime = unit.toNanos(keepAliveTime);
		this.threadFactory = threadFactory;
		this.handler = handler;
	}

	public void execute(Runnable command) {
		if (command == null)
			throw new NullPointerException();
		taskCount.incrementAndGet();
		if (poolSize >= corePoolSize || !addIfUnderCorePoolSize(command)) {
			if (runState == RUNNING && workQueue.offer(command)) {
				if (runState != RUNNING || waitingWorker.get() == 0) {
					ensureQueuedTaskHandled(command);
				}
			} else if (!addIfUnderMaximumPoolSize(command)) {
				reject(command); // is shutdown or saturated
			}
		}
	}

	private Thread addThread(Runnable firstTask) {
		Worker w = new Worker(firstTask);
		Thread t = threadFactory.newThread(w);
		if (t != null) {
			w.thread = t;
			workers.add(w);
			int nt = ++poolSize;
			if (nt > largestPoolSize)
				largestPoolSize = nt;
		}
		return t;
	}

	private boolean addIfUnderCorePoolSize(Runnable firstTask) {
		Thread t = null;
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			if (poolSize < corePoolSize && runState == RUNNING)
				t = addThread(firstTask);
		} finally {
			mainLock.unlock();
		}
		if (t == null)
			return false;
		t.start();
		return true;
	}

	private boolean addIfUnderMaximumPoolSize(Runnable firstTask) {
		Thread t = null;
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			if (poolSize < maximumPoolSize && runState == RUNNING)
				t = addThread(firstTask);
		} finally {
			mainLock.unlock();
		}
		if (t == null)
			return false;
		t.start();
		return true;
	}

	private void ensureQueuedTaskHandled(Runnable command) {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		boolean reject = false;
		Thread t = null;
		try {
			int state = runState;
			if (state != RUNNING && workQueue.remove(command))
				reject = true;
			else if (state < STOP
					&& poolSize < Math.max(this.maximumPoolSize, 1)
					&& !workQueue.isEmpty())
				t = addThread(null);
		} finally {
			mainLock.unlock();
		}
		if (reject)
			reject(command);
		else if (t != null)
			t.start();
	}

	void reject(Runnable command) {
		handler.rejectedExecution(command, this);
	}

	private final class Worker implements Runnable {

		private final ReentrantLock runLock = new ReentrantLock();

		private Runnable firstTask;

		volatile long completedTasks;

		Thread thread;

		Worker(Runnable firstTask) {
			this.firstTask = firstTask;
		}

		boolean isActive() {
			return runLock.isLocked();
		}

		void interruptIfIdle() {
			final ReentrantLock runLock = this.runLock;
			if (runLock.tryLock()) {
				try {
					if (thread != Thread.currentThread())
						thread.interrupt();
				} finally {
					runLock.unlock();
				}
			}
		}

		void interruptNow() {
			thread.interrupt();
		}

		private void runTask(Runnable task) {
			final ReentrantLock runLock = this.runLock;
			runLock.lock();
			try {
				if (runState < STOP && Thread.interrupted() && runState >= STOP)
					thread.interrupt();

				boolean ran = false;
				beforeExecute(thread, task);
				try {
					task.run();
					completedTaskCount.incrementAndGet();
					ran = true;
					afterExecute(task, null);
					++completedTasks;
				} catch (RuntimeException ex) {
					if (!ran)
						afterExecute(task, ex);
					throw ex;
				}
			} finally {
				runLock.unlock();
			}
		}

		public void run() {
			log.info("thread worker start ({}/{},{})", new Object[] { poolSize,
					corePoolSize, maximumPoolSize });
			try {
				Runnable task = firstTask;
				firstTask = null;
				while (task != null || (task = getTask()) != null) {
					runTask(task);
					task = null;
				}
			} finally {
				workerDone(this);
				log.info("thread worker exit  ({}/{},{})", new Object[] {
						poolSize, corePoolSize, maximumPoolSize });
			}
		}
	}

	Runnable getTask() {
		for (;;) {
			try {
				int state = runState;
				if (state > SHUTDOWN)
					return null;
				Runnable r;
				waitingWorker.incrementAndGet();
				try {
					if (state == SHUTDOWN) // Help drain queue
						r = workQueue.poll();
					else if (poolSize > corePoolSize || allowCoreThreadTimeOut)
						r = workQueue.poll(keepAliveTime, TimeUnit.NANOSECONDS);
					else
						r = workQueue.take();
				} finally {
					waitingWorker.decrementAndGet();
				}
				if (r != null)
					return r;
				if (workerCanExit()) {
					if (runState >= SHUTDOWN) // Wake up others
						interruptIdleWorkers();
					return null;
				}
				// Else retry
			} catch (InterruptedException ie) {
				// On interruption, re-check runState
			}
		}
	}

	private boolean workerCanExit() {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		boolean canExit;
		try {
			canExit = runState >= STOP
					|| (runState == SHUTDOWN && workQueue.isEmpty())
					|| allowCoreThreadTimeOut
					|| poolSize > Math.max(1, corePoolSize);
		} finally {
			mainLock.unlock();
		}
		return canExit;
	}

	void interruptIdleWorkers() {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			for (Worker w : workers)
				w.interruptIfIdle();
		} finally {
			mainLock.unlock();
		}
	}

	void workerDone(Worker w) {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			workers.remove(w);
			if (--poolSize == 0)
				tryTerminate();
		} finally {
			mainLock.unlock();
		}
	}

	/* Termination support. */
	private void tryTerminate() {
		if (poolSize == 0) {
			int state = runState;
			if (state < STOP && !workQueue.isEmpty()) {
				state = RUNNING; // disable termination check below
				Thread t = addThread(null);
				if (t != null)
					t.start();
			}
			if (state == STOP || state == SHUTDOWN) {
				runState = TERMINATED;
				termination.signalAll();
				terminated();
			}
		}
	}

	public void shutdown() {
		SecurityManager security = System.getSecurityManager();
		if (security != null)
			security.checkPermission(shutdownPerm);

		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			if (security != null) { // Check if caller can modify our threads
				for (Worker w : workers)
					security.checkAccess(w.thread);
			}

			int state = runState;
			if (state < SHUTDOWN)
				runState = SHUTDOWN;

			try {
				for (Worker w : workers) {
					w.interruptIfIdle();
				}
			} catch (SecurityException se) { // Try to back out
				runState = state;
				// tryTerminate() here would be a no-op
				throw se;
			}

			tryTerminate(); // Terminate now if pool and queue empty
		} finally {
			mainLock.unlock();
		}
	}

	public List<Runnable> shutdownNow() {
		SecurityManager security = System.getSecurityManager();
		if (security != null)
			security.checkPermission(shutdownPerm);

		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			if (security != null) { // Check if caller can modify our threads
				for (Worker w : workers)
					security.checkAccess(w.thread);
			}

			int state = runState;
			if (state < STOP)
				runState = STOP;

			try {
				for (Worker w : workers) {
					w.interruptNow();
				}
			} catch (SecurityException se) { // Try to back out
				runState = state;
				// tryTerminate() here would be a no-op
				throw se;
			}

			List<Runnable> tasks = drainQueue();
			tryTerminate(); // Terminate now if pool and queue empty
			return tasks;
		} finally {
			mainLock.unlock();
		}
	}

	private List<Runnable> drainQueue() {
		List<Runnable> taskList = new ArrayList<Runnable>();
		workQueue.drainTo(taskList);
		/*
		 * If the queue is a DelayQueue or any other kind of queue for which
		 * poll or drainTo may fail to remove some elements, we need to manually
		 * traverse and remove remaining tasks. To guarantee atomicity wrt other
		 * threads using this queue, we need to create a new iterator for each
		 * element removed.
		 */
		while (!workQueue.isEmpty()) {
			Iterator<Runnable> it = workQueue.iterator();
			try {
				if (it.hasNext()) {
					Runnable r = it.next();
					if (workQueue.remove(r))
						taskList.add(r);
				}
			} catch (ConcurrentModificationException ignore) {
			}
		}
		return taskList;
	}

	public boolean isShutdown() {
		return runState != RUNNING;
	}

	public boolean isTerminating() {
		int state = runState;
		return state == SHUTDOWN || state == STOP;
	}

	public boolean isTerminated() {
		return runState == TERMINATED;
	}

	public boolean awaitTermination(long timeout, TimeUnit unit)
			throws InterruptedException {
		long nanos = unit.toNanos(timeout);
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			for (;;) {
				if (runState == TERMINATED)
					return true;
				if (nanos <= 0)
					return false;
				nanos = termination.awaitNanos(nanos);
			}
		} finally {
			mainLock.unlock();
		}
	}

	protected void finalize() {
		shutdown();
	}

	/* Getting and setting tunable parameters */
	public void setThreadFactory(ThreadFactory threadFactory) {
		if (threadFactory == null)
			throw new NullPointerException();
		this.threadFactory = threadFactory;
	}

	public ThreadFactory getThreadFactory() {
		return threadFactory;
	}

	public void setRejectedExecutionHandler(BMARejectedExecutionHandler handler) {
		if (handler == null)
			throw new NullPointerException();
		this.handler = handler;
	}

	public BMARejectedExecutionHandler getRejectedExecutionHandler() {
		return handler;
	}

	public void setCorePoolSize(int corePoolSize) {
		if (corePoolSize < 0)
			throw new IllegalArgumentException();
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			int extra = this.corePoolSize - corePoolSize;
			this.corePoolSize = corePoolSize;
			if (extra < 0) {
				int n = workQueue.size(); // don't add more threads than tasks
				while (extra++ < 0 && n-- > 0 && poolSize < corePoolSize) {
					Thread t = addThread(null);
					if (t != null)
						t.start();
					else
						break;
				}
			} else if (extra > 0 && poolSize > corePoolSize) {
				try {
					Iterator<Worker> it = workers.iterator();
					while (it.hasNext() && extra-- > 0
							&& poolSize > corePoolSize
							&& workQueue.remainingCapacity() == 0)
						it.next().interruptIfIdle();
				} catch (SecurityException ignore) {
					// Not an error; it is OK if the threads stay live
				}
			}
		} finally {
			mainLock.unlock();
		}
	}

	public int getCorePoolSize() {
		return corePoolSize;
	}

	public int getRunState() {
		return runState;
	}

	public boolean prestartCoreThread() {
		return addIfUnderCorePoolSize(null);
	}

	public int prestartAllCoreThreads() {
		int n = 0;
		while (addIfUnderCorePoolSize(null))
			++n;
		return n;
	}

	public boolean allowsCoreThreadTimeOut() {
		return allowCoreThreadTimeOut;
	}

	public void allowCoreThreadTimeOut(boolean value) {
		if (value && keepAliveTime <= 0)
			throw new IllegalArgumentException(
					"Core threads must have nonzero keep alive times");

		allowCoreThreadTimeOut = value;
	}

	public void setMaximumPoolSize(int maximumPoolSize) {
		if (maximumPoolSize <= 0 || maximumPoolSize < corePoolSize)
			throw new IllegalArgumentException();
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			int extra = this.maximumPoolSize - maximumPoolSize;
			this.maximumPoolSize = maximumPoolSize;
			if (extra > 0 && poolSize > maximumPoolSize) {
				try {
					Iterator<Worker> it = workers.iterator();
					while (it.hasNext() && extra > 0
							&& poolSize > maximumPoolSize) {
						it.next().interruptIfIdle();
						--extra;
					}
				} catch (SecurityException ignore) {
					// Not an error; it is OK if the threads stay live
				}
			}
		} finally {
			mainLock.unlock();
		}
	}

	public int getMaximumPoolSize() {
		return maximumPoolSize;
	}

	public void setKeepAliveTime(long time, TimeUnit unit) {
		if (time < 0)
			throw new IllegalArgumentException();
		if (time == 0 && allowsCoreThreadTimeOut())
			throw new IllegalArgumentException(
					"Core threads must have nonzero keep alive times");
		this.keepAliveTime = unit.toNanos(time);
	}

	public long getKeepAliveTime(TimeUnit unit) {
		return unit.convert(keepAliveTime, TimeUnit.NANOSECONDS);
	}

	/* User-level queue utilities */
	public BlockingQueue<Runnable> getQueue() {
		return workQueue;
	}

	public boolean remove(Runnable task) {
		return getQueue().remove(task);
	}

	public void purge() {
		// Fail if we encounter interference during traversal
		try {
			Iterator<Runnable> it = getQueue().iterator();
			while (it.hasNext()) {
				Runnable r = it.next();
				if (r instanceof Future<?>) {
					Future<?> c = (Future<?>) r;
					if (c.isCancelled())
						it.remove();
				}
			}
		} catch (ConcurrentModificationException ex) {
			return;
		}
	}

	/* Statistics */
	public int getPoolSize() {
		return poolSize;
	}

	public int getActiveCount() {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			int n = 0;
			for (Worker w : workers) {
				if (w.isActive())
					++n;
			}
			return n;
		} finally {
			mainLock.unlock();
		}
	}

	public int getLargestPoolSize() {
		final ReentrantLock mainLock = this.mainLock;
		mainLock.lock();
		try {
			return largestPoolSize;
		} finally {
			mainLock.unlock();
		}
	}

	public long getTaskCount() {
		return taskCount.get();
	}

	public long getCompletedTaskCount() {
		return completedTaskCount.get();
	}

	/* Extension hooks */
	protected void beforeExecute(Thread t, Runnable r) {
	}

	protected void afterExecute(Runnable r, Throwable t) {
	}

	protected void terminated() {
		ThreadPoolSquare.remove(this);
	}

	/**
	 * A handler for rejected tasks that runs the rejected task directly in the
	 * calling thread of the <tt>execute</tt> method, unless the executor has
	 * been shut down, in which case the task is discarded.
	 */
	public static class CallerRunsPolicy implements BMARejectedExecutionHandler {
		/**
		 * Creates a <tt>CallerRunsPolicy</tt>.
		 */
		public CallerRunsPolicy() {
		}

		/**
		 * Executes task r in the caller's thread, unless the executor has been
		 * shut down, in which case the task is discarded.
		 * 
		 * @param r
		 *            the runnable task requested to be executed
		 * @param e
		 *            the executor attempting to execute this task
		 */
		public void rejectedExecution(Runnable r, BMAThreadPoolExecutor e) {
			if (!e.isShutdown()) {
				r.run();
			}
		}
	}

	/**
	 * A handler for rejected tasks that throws a
	 * <tt>RejectedExecutionException</tt>.
	 */
	public static class AbortPolicy implements BMARejectedExecutionHandler {
		/**
		 * Creates an <tt>AbortPolicy</tt>.
		 */
		public AbortPolicy() {
		}

		/**
		 * Always throws RejectedExecutionException.
		 * 
		 * @param r
		 *            the runnable task requested to be executed
		 * @param e
		 *            the executor attempting to execute this task
		 * @throws RejectedExecutionException
		 *             always.
		 */
		public void rejectedExecution(Runnable r, BMAThreadPoolExecutor e) {
			throw new RejectedExecutionException();
		}
	}

	/**
	 * A handler for rejected tasks that silently discards the rejected task.
	 */
	public static class DiscardPolicy implements BMARejectedExecutionHandler {
		/**
		 * Creates a <tt>DiscardPolicy</tt>.
		 */
		public DiscardPolicy() {
		}

		/**
		 * Does nothing, which has the effect of discarding task r.
		 * 
		 * @param r
		 *            the runnable task requested to be executed
		 * @param e
		 *            the executor attempting to execute this task
		 */
		public void rejectedExecution(Runnable r, BMAThreadPoolExecutor e) {
		}
	}

	/**
	 * A handler for rejected tasks that discards the oldest unhandled request
	 * and then retries <tt>execute</tt>, unless the executor is shut down, in
	 * which case the task is discarded.
	 */
	public static class DiscardOldestPolicy implements
			BMARejectedExecutionHandler {
		/**
		 * Creates a <tt>DiscardOldestPolicy</tt> for the given executor.
		 */
		public DiscardOldestPolicy() {
		}

		/**
		 * Obtains and ignores the next task that the executor would otherwise
		 * execute, if one is immediately available, and then retries execution
		 * of task r, unless the executor is shut down, in which case task r is
		 * instead discarded.
		 * 
		 * @param r
		 *            the runnable task requested to be executed
		 * @param e
		 *            the executor attempting to execute this task
		 */
		public void rejectedExecution(Runnable r, BMAThreadPoolExecutor e) {
			if (!e.isShutdown()) {
				e.getQueue().poll();
				e.execute(r);
			}
		}
	}
}
